// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
// fork from CAFFE2 registry

#pragma once

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <memory>
#include <mutex>
#include <string_view>
#include <unordered_map>

#include <turbo/log/logging.h>
#include "melon/preprocessor.h"
#include <melon/container/f14_map.h>
#include <pollux/common/base/exceptions.h>
#include <pollux/common/meta_programming.h>

namespace kumo {
    namespace pollux {
        /**
         * @brief A template class that allows one to register function objects by keys.
         *
         * The keys are usually a string specifying the name, but can be anything that
         * can be used in a std::map.
         * It provides `Create` function to directly call with key and arguments.
         */

        template<class KeyType, class FunctionSignature>
        class Registry {
        public:
            using Creator = std::function<FunctionSignature>;
            using CreatorMap = melon::F14NodeMap<KeyType, Creator>;

            Registry() : Create(creatorMap_) {
            }

            void Register(
                const KeyType &key,
                Creator creator,
                std::optional<std::string_view> helpMsg = std::nullopt) {
                std::lock_guard<std::mutex> lock(registerutex_);
                creatorMap_[key] = std::move(creator);
                if (helpMsg) {
                    helpMessage_[key] = *helpMsg;
                }
            }

            inline bool Has(const KeyType &key) const {
                return (creatorMap_.count(key) != 0);
            }

            /**
             * Returns the keys currently registered as a vector.
             */
            std::vector<KeyType> Keys() const {
                std::vector<KeyType> keys;
                for (const auto &it: creatorMap_) {
                    keys.push_back(it.first);
                }
                return keys;
            }

            const melon::F14FastMap<KeyType, std::string> &HelpMessage() const {
                return helpMessage_;
            }

            const char *HelpMessage(const KeyType &key) const {
                auto it = helpMessage_.find(key);
                if (it == helpMessage_.end()) {
                    return nullptr;
                }
                return it->second.c_str();
            }

        private:
            CreatorMap creatorMap_;
            melon::F14FastMap<KeyType, std::string> helpMessage_;
            std::mutex registerutex_;

            Registry(const Registry &other) = delete;

        public:
            template<typename T>
            struct CreateFunction;

            template<class ReturnType, class... ArgTypes>
            struct CreateFunction<ReturnType(ArgTypes...)> {
                CreateFunction(const CreatorMap &creatorMap) : creatorMap_(creatorMap) {
                }

                ReturnType operator()(const KeyType &key, ArgTypes... types) const {
                    const auto it = creatorMap_.find(key);
                    if (it == creatorMap_.end()) {
                        if constexpr (kumo::pollux::util::is_smart_pointer<
                            ReturnType>::value) {
                            return nullptr;
                        }

                        POLLUX_UNSUPPORTED(
                            typeid(ReturnType).name(), " is not nullable return type");
                    }
                    return it->second(types...);
                }

            private:
                const CreatorMap &creatorMap_;
            };

            // Provide Create function as a function object
            // If there is no key found, it will:
            //   Smart pointer types: return nullptr
            //   Value types        : throw invalid_argument exception.
            //
            // Function signature:
            //   ReturnType Create(const KeyType& key, ArgTypes...);
            const CreateFunction<FunctionSignature> Create;
        };
    } // namespace pollux
} // namespace kumo
